#include "mm.h"
#include "system.h"


struct bucket_desc
{
    void *page;
    struct bucket_desc *next;
    void *freeptr;
    u16 refcnt;
    u16 bucket_size;
} __attribute__((packed));

struct _bucket_dir
{
    int size;
    struct bucket_desc *chain;
} __attribute__((packed));

struct _bucket_dir bucket_dir[] =
{
    {16, (struct bucket_desc *) NULL},
    {32, (struct bucket_desc *) NULL},
    {64, (struct bucket_desc *) NULL},
    {128, (struct bucket_desc *) NULL},
    {256, (struct bucket_desc *) NULL},
    {512, (struct bucket_desc *) NULL},
    {1024, (struct bucket_desc *) NULL},
    {2048, (struct bucket_desc *) NULL},
    {4096, (struct bucket_desc *) NULL},
    {0, NULL},
};

struct bucket_desc *free_bucket_desc = (struct bucket_desc *) NULL;

static inline void init_bucket_desc()
{
    struct bucket_desc *bdesc, *first;
    int i;

    first = bdesc = (struct bucket_desc *) get_free_page();
    if (!bdesc)
    {
        return ;
    }
    for (i = PAGE_SIZE / sizeof(struct bucket_desc); i > 1; i--)
    {
        bdesc->next = bdesc + 1;
        bdesc++;
    }
    
    bdesc->next = free_bucket_desc;
    free_bucket_desc = first;
}
/**
 * @brief 该函数是一个内存分配函数，它的目的是在内部维护的链表中找到合适大小的空闲块，将其分配给调用者并返回分配的内存指针。函数的输入参数为需要分配的内存大小，返回值为分配的内存指针。

函数首先通过遍历链表中的桶（bucket）来找到合适的大小，其中桶是一个结构体，包含块的大小和链表头指针。如果找到了合适大小的桶，则从桶的链表中查找是否有可用的空闲块。如果没有找到可用的空闲块，则需要创建一个新的空闲块。

如果遍历所有桶后仍未找到可用的空闲块，说明请求的内存大小超出了已分配的内存大小，则函数会输出错误信息并返回NULL。否则，函数会在内部维护的链表中创建一个新的空闲块来满足请求。具体步骤如下：

首先，函数定义了一些局部变量，包括桶指针bdir、桶描述符指针bdesc和返回值指针retval。
然后，函数使用一个循环来遍历链表中的桶，直到找到大小不小于请求大小的桶，或者遍历完所有桶。如果遍历完所有桶都未找到合适大小的桶，则函数会输出错误信息并返回NULL。
在找到合适的桶之后，函数使用另一个循环来在桶的链表中查找是否有可用的空闲块。如果没有找到可用的空闲块，则需要创建一个新的空闲块。
如果在内部维护的链表中没有可用的空闲块，函数会检查是否还有未使用的桶描述符。如果有，则取一个出来使用；如果没有，则调用另一个函数来创建一个新的桶描述符。
获取到桶描述符后，函数会根据桶的大小，使用get_free_page()函数来获取一个自由页（即一块连续的内存），并将该页的地址赋给桶描述符的page和freeptr成员。然后，函数根据桶的大小，使用指针的类型转换将自由页地址转换为指针类型。
接下来，函数使用循环在该自由页中创建多个内存块，每个内存块的大小等于桶的大小。具体做法是将每个内存块的地址赋给一个指向下一个块的指针，并将指针自身存储在类型为指针类型的数组中。最后，将数组的最后一个元素设置为NULL，表示链表结束。
创建完内存块后，将桶描述符的next成员设置为之前桶的链表头指针，将桶链表头指针设置为新的桶描述符，将新的桶描述符设置为可用状态。
最后，将桶描述符的freeptr成员赋给retval，将桶描述符的refcnt成员增加1，然后将retval作为函数的返回值返回给调用者。
 * 
 * @param len 
 * @return void* 
 */
void* kmalloc(size_t len)
{
    struct _bucket_dir *bdir;
    struct bucket_desc *bdesc;
    void *retval;

    for (bdir = bucket_dir; bdir->size; bdir++)
    {
        if (bdir->size >= len)
        {
            break;
        }
    }
    
    if (!bdir->size)
    {
        printk("kmalloc called with impossibly large argument (%d)\n", len);
        return NULL;
    }
    
    for (bdesc = bdir->chain; bdesc; bdesc = bdesc->next)
    {
        if (bdesc->freeptr)
        {
            break;
        }
    }

    if (!bdesc)
    {
        char *cp;
        int i;
        if (!free_bucket_desc)
        {
            init_bucket_desc();
        }
        bdesc = free_bucket_desc;
        free_bucket_desc = bdesc->next;
        bdesc->refcnt = 0;
        bdesc->bucket_size = bdir->size;
        bdesc->page = bdesc->freeptr = (void*)(cp = (char*)get_free_page());
        if (!cp)
        {
            return NULL;
        }
        for (i = PAGE_SIZE / bdir->size; i > 1; i--)
        {
            *((char **) cp) = cp + bdir->size;
            cp += bdir->size;
        }
        *((char **) cp) = NULL;
        bdesc->next = bdir->chain;
        bdir->chain = bdesc;
    }
    retval = bdesc->freeptr;
    bdesc->freeptr = *((void **) retval);
    bdesc->refcnt ++;
    return (retval);
}

void kfree_s(void* obj, int size)
{
    void *page;
    struct _bucket_dir *bdir;
    struct bucket_desc *bdesc, *prev;
    bdesc = prev = 0;

    page = (void*) ((unsigned long) obj & 0xfffff000);
    for (bdir = bucket_dir; bdir->size; bdir++)
    {
        prev = 0;
        if (bdir->size < size)
        {
            continue;
        }
        for (bdir = bucket_dir; bdir->size; bdir++)
        {
            if (bdesc->page == page)
            {
                goto found;
            }
            prev = bdesc;
        }
    }
    return;
found:
    *((void**)obj) = bdesc->freeptr;
    bdesc->freeptr = obj;
    bdesc->refcnt--;

    if (bdesc->refcnt == 0)
    {
        if ((prev && (prev->next != bdesc)) || 
          (!prev && (bdir->chain != bdesc)))
        {
            for (prev = bdir->chain; prev; prev = prev->next)
            {
                if (prev->next == bdesc)
                {
                    break;
                }
            }
        }
        if (prev)   
        {
            prev->next = bdesc->next;
        }
        else
        {
            if (bdir->chain != bdesc)
            {
                return;
            }
            bdir->chain = bdesc->next;
        }
        free_page((unsigned long) bdesc->page);
        bdesc->next = free_bucket_desc;
        free_bucket_desc = bdesc;
    }
    return;
}